package ch.ethz.fcl.tools;

import java.awt.Component;
import java.awt.Dimension;
import java.awt.Graphics;
import java.awt.event.KeyAdapter;
import java.awt.event.KeyEvent;
import java.awt.event.WindowAdapter;
import java.awt.event.WindowEvent;
import java.awt.geom.AffineTransform;
import java.awt.image.AffineTransformOp;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
import java.awt.image.DataBufferInt;
import java.awt.image.DataBufferUShort;
import java.awt.image.Raster;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

import javax.imageio.ImageIO;
import javax.swing.JFrame;

import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.core.MatOfPoint;
import org.opencv.core.Point;
import org.opencv.core.Scalar;
import org.opencv.highgui.Highgui;
import org.opencv.imgproc.Imgproc;

import ch.ethz.fcl.Constant;

public class ImageTools {
	public static BufferedImage matToImage(Mat mat) {
		MatOfByte matOfByte = new MatOfByte();

		Highgui.imencode(".png", mat, matOfByte);

		byte[] byteArray = matOfByte.toArray();
		BufferedImage bufImage = null;

		try {
			InputStream in = new ByteArrayInputStream(byteArray);
			bufImage = ImageIO.read(in);
		} catch (Exception e) {
			e.printStackTrace();
		}

		// Flip the image vertically
		AffineTransform tx = AffineTransform.getScaleInstance(1, -1);
		tx.translate(0, -bufImage.getHeight(null));
		AffineTransformOp op = new AffineTransformOp(tx,
				AffineTransformOp.TYPE_NEAREST_NEIGHBOR);
		bufImage = op.filter(bufImage, null);

		return bufImage;
	}

	public static Mat imageToMat(BufferedImage image) {
		Mat mat = null;

		if (image.getType() == BufferedImage.TYPE_BYTE_GRAY) {
			mat = new Mat(image.getHeight(), image.getWidth(), CvType.CV_8UC1);

			byte[] data = ((DataBufferByte) image.getRaster().getDataBuffer())
					.getData();
			mat.put(0, 0, data);
		} else if (image.getType() == BufferedImage.TYPE_USHORT_GRAY) {
			mat = new Mat(image.getHeight(), image.getWidth(), CvType.CV_16UC1);

			short[] data = ((DataBufferUShort) image.getRaster()
					.getDataBuffer()).getData();

			// System.out.println("image width "+image.getWidth()+" height "+image.getHeight());
			// System.out.println("mat size "+mat.size());
			mat.put(0, 0, data);
		} else if (image.getType() == BufferedImage.TYPE_INT_RGB) {
			mat = new Mat(image.getHeight(), image.getWidth(), CvType.CV_8UC3);

			int[] pixels = ((DataBufferInt) image.getRaster().getDataBuffer())
					.getData();
			byte[] pixelBytes = new byte[pixels.length * 3];
			for (int i = 0; i < pixels.length; i++) {
				pixelBytes[i * 3] = (byte) ((pixels[i] >> 16) & 0xFF);
				pixelBytes[i * 3 + 1] = (byte) ((pixels[i] >> 8) & 0xFF);
				pixelBytes[i * 3] = (byte) ((pixels[i]) & 0xFF);
			}

			mat.put(0, 0, pixelBytes);
		}

		return mat;
	}

	public static int[] imageToArray(BufferedImage image) {
		int[] array = new int[image.getWidth() * image.getHeight()];
		image.getRGB(0, 0, image.getWidth(), image.getHeight(), array, 0,
				image.getWidth());
		return array;
	}

	public static BufferedImage arrayToImage(int[] array, int width, int height) {
		BufferedImage image = new BufferedImage(width, height,
				BufferedImage.TYPE_INT_RGB);
		image.setRGB(0, 0, width, height, array, 0, width);
		return image;
	}

	public static BufferedImage arrayToImage(byte[] array, int width, int height) {
		BufferedImage image = new BufferedImage(width, height,
				BufferedImage.TYPE_BYTE_GRAY);

		DataBufferByte buffer = new DataBufferByte(array, width * height);
		Raster raster = Raster.createPackedRaster(buffer, width, height, 8,
				null);
		image.setData(raster);
		return image;
	}

	public static void writeImage(BufferedImage image, String file) {
		try {
			File outputfile = new File(file);
			ImageIO.write(image, "png", outputfile);
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public static void draw(Mat mat) {
		draw(matToImage(mat));
	}

	public static void draw(BufferedImage img) {
		JFrame frame = new JFrame("test");

		frame.addKeyListener(new KeyAdapter() {
			@Override
			public void keyPressed(KeyEvent e) {
				if (e.getKeyCode() == KeyEvent.VK_ESCAPE) {
					System.exit(0);
				}
			}
		});

		frame.addWindowListener(new WindowAdapter() {
			@Override
			public void windowClosing(WindowEvent e) {
				System.exit(0);
			}
		});

		frame.setLocation(500, 300);
		frame.add("Center", new ImageTools.ImageIcon(img));
		frame.pack();
		frame.setVisible(true);
	}

	// the function draws all the polys in the image
	public static Mat drawPolys(Mat image, List<MatOfPoint> allPolys) {
		Mat output = new Mat(image.size(), CvType.CV_8UC3);
		if (image.channels() == 1)
			Imgproc.cvtColor(image, output, Imgproc.COLOR_GRAY2RGB);
		else if (image.channels() == 3)
			image.copyTo(output);

		for (int i = 0; i < allPolys.size(); i++) {
			MatOfPoint i_poly = allPolys.get(i);
			if (i_poly == null)
				continue;

			List<MatOfPoint> p = new ArrayList<MatOfPoint>();
			p.add(i_poly);

			Scalar color = ImageTools.getColor(i);
			Core.polylines(output, p, true, color, 1, Core.LINE_AA, 0);
		}
		return output;
	}

	// angle (in radians) between p0->p1 and p2->p3
	public static double angle(Point p0, Point p1, Point p2, Point p3) {
		double dx1 = p1.x - p0.x;
		double dy1 = p1.y - p0.y;
		double dx2 = p3.x - p2.x;
		double dy2 = p3.y - p2.y;

		double m1 = Math.sqrt(dx1 * dx1 + dy1 * dy1);
		double m2 = Math.sqrt(dx2 * dx2 + dy2 * dy2);

		return Math.acos((dx1 * dx2 + dy1 * dy2) / (m1 * m2));
	}

	public static Scalar getColor(int index) {
		switch (index) {
		case 0:
			return new Scalar(255 * Constant.COLOR_SCALE,
					100 * Constant.COLOR_SCALE, 200 * Constant.COLOR_SCALE);
		case 1:
			return new Scalar(255 * Constant.COLOR_SCALE, 0, 0); // blue
		case 2:
			return new Scalar(0, 255 * Constant.COLOR_SCALE, 0); // green
		case 3:
			return new Scalar(0, 0, 255 * Constant.COLOR_SCALE); // red
		case 4:
			return new Scalar(255 * Constant.COLOR_SCALE,
					255 * Constant.COLOR_SCALE, 0);
		case 5:
			return new Scalar(255 * Constant.COLOR_SCALE, 0,
					255 * Constant.COLOR_SCALE);
		case 6:
			return new Scalar(100 * Constant.COLOR_SCALE,
					100 * Constant.COLOR_SCALE, 255 * Constant.COLOR_SCALE);
		case 7:
			return new Scalar(100 * Constant.COLOR_SCALE,
					255 * Constant.COLOR_SCALE, 255 * Constant.COLOR_SCALE);
		case 8:
			return new Scalar(255 * Constant.COLOR_SCALE,
					155 * Constant.COLOR_SCALE, 100 * Constant.COLOR_SCALE);
		default:
			return new Scalar(0, 100 * Constant.COLOR_SCALE,
					255 * Constant.COLOR_SCALE);
		}
	}

	private static class ImageIcon extends Component {
		private static final long serialVersionUID = 3555287429032087653L;
		private BufferedImage image;
		private int width, height;

		public ImageIcon(BufferedImage image) {
			this.image = image;
			this.width = image.getWidth();
			this.height = image.getHeight();
		}

		@Override
		public Dimension getPreferredSize() {
			return new Dimension(width, height);
		}

		@Override
		public void paint(Graphics g) {
			g.drawImage(image, 0, 0, null);
		}
	}
}
